{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "826IBSWMN4rr"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ITj3u97-tNR7"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BYwfpc4wN4rt"
      },
      "source": [
        "# 权重聚类综合指南"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IFva_Ed5N4ru"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://tensorflow.google.cn/model_optimization/guide/clustering/clustering_comprehensive_guide\">     <img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">     在 TensorFlow.org 上查看</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/model_optimization/guide/clustering/clustering_comprehensive_guide.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\">在 Google Colab 中运行 </a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/model_optimization/guide/clustering/clustering_comprehensive_guide.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">在 GitHub 上查看源代码</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/model_optimization/guide/clustering/clustering_comprehensive_guide.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\"> 下载笔记本</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tidmcl3sN4rv"
      },
      "source": [
        "欢迎阅读 TensorFlow Model Optimization Toolkit 中*权重聚类*的综合指南。\n",
        "\n",
        "本页面记录了各种用例，并展示了如何将 API 用于每种用例​​。了解需要哪些 API 后，可在 [API 文档](https://tensorflow.google.cn/model_optimization/api_docs/python/tfmot/clustering)中找到参数和底层详细信息：\n",
        "\n",
        "- 如果要查看权重聚类的好处以及支持的功能，请查看[概述](https://tensorflow.google.cn/model_optimization/guide/clustering)。\n",
        "- 有关单个端到端示例，请参阅[权重聚类示例](https://tensorflow.google.cn/model_optimization/guide/clustering/clustering_example)。\n",
        "\n",
        "本指南涵盖了以下用例：\n",
        "\n",
        "- 定义聚类模型。\n",
        "- 为聚类模型设置检查点和进行反序列化。\n",
        "- 提高聚类模型的准确率。\n",
        "- 仅对于部署而言，您必须采取措施才能看到压缩的好处。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RRtKxbo8N4rv"
      },
      "source": [
        "## 设置\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "08dJRvOqN4rw"
      },
      "outputs": [],
      "source": [
        "! pip install -q tensorflow-model-optimization\n",
        "\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import tempfile\n",
        "import os\n",
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "input_dim = 20\n",
        "output_dim = 20\n",
        "x_train = np.random.randn(1, input_dim).astype(np.float32)\n",
        "y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)\n",
        "\n",
        "def setup_model():\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(input_dim, input_shape=[input_dim]),\n",
        "      tf.keras.layers.Flatten()\n",
        "  ])\n",
        "  return model\n",
        "\n",
        "def train_model(model):\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.categorical_crossentropy,\n",
        "      optimizer='adam',\n",
        "      metrics=['accuracy']\n",
        "  )\n",
        "  model.summary()\n",
        "  model.fit(x_train, y_train)\n",
        "  return model\n",
        "\n",
        "def save_model_weights(model):\n",
        "  _, pretrained_weights = tempfile.mkstemp('.h5')\n",
        "  model.save_weights(pretrained_weights)\n",
        "  return pretrained_weights\n",
        "\n",
        "def setup_pretrained_weights():\n",
        "  model= setup_model()\n",
        "  model = train_model(model)\n",
        "  pretrained_weights = save_model_weights(model)\n",
        "  return pretrained_weights\n",
        "\n",
        "def setup_pretrained_model():\n",
        "  model = setup_model()\n",
        "  pretrained_weights = setup_pretrained_weights()\n",
        "  model.load_weights(pretrained_weights)\n",
        "  return model\n",
        "\n",
        "def save_model_file(model):\n",
        "  _, keras_file = tempfile.mkstemp('.h5') \n",
        "  model.save(keras_file, include_optimizer=False)\n",
        "  return keras_file\n",
        "\n",
        "def get_gzipped_model_size(model):\n",
        "  # It returns the size of the gzipped model in bytes.\n",
        "  import os\n",
        "  import zipfile\n",
        "\n",
        "  keras_file = save_model_file(model)\n",
        "\n",
        "  _, zipped_file = tempfile.mkstemp('.zip')\n",
        "  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:\n",
        "    f.write(keras_file)\n",
        "  return os.path.getsize(zipped_file)\n",
        "\n",
        "setup_model()\n",
        "pretrained_weights = setup_pretrained_weights()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ARd37qONN4rz"
      },
      "source": [
        "## 定义聚类模型\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHB3pkU3N4r0"
      },
      "source": [
        "### 聚类整个模型（序贯模型和函数式模型）"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ig-il1lmN4r1"
      },
      "source": [
        "提高模型准确率的**提示**：\n",
        "\n",
        "- 您必须将具有可接受准确率的预训练模型传递给此 API。使用聚类从头开始训练模型会导致准确率不佳。\n",
        "- 在某些情况下，聚类某些层会对模型准确率造成不利影响。查看“聚类某些层”来了解如何跳过聚类对准确率影响最大的层。\n",
        "\n",
        "要聚类所有层，请将 `tfmot.clustering.keras.cluster_weights` 应用于模型。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29g7OADjN4r1"
      },
      "outputs": [],
      "source": [
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "cluster_weights = tfmot.clustering.keras.cluster_weights\n",
        "CentroidInitialization = tfmot.clustering.keras.CentroidInitialization\n",
        "\n",
        "clustering_params = {\n",
        "  'number_of_clusters': 3,\n",
        "  'cluster_centroids_init': CentroidInitialization.DENSITY_BASED\n",
        "}\n",
        "\n",
        "model = setup_model()\n",
        "model.load_weights(pretrained_weights)\n",
        "\n",
        "clustered_model = cluster_weights(model, **clustering_params)\n",
        "\n",
        "clustered_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zEOHK4OON4r7"
      },
      "source": [
        "### 聚类某些层（序贯模型和函数式模型）\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENscQ7ZWN4r8"
      },
      "source": [
        "提高模型准确率的**提示**：\n",
        "\n",
        "- 您必须将具有可接受准确率的预训练模型传递给此 API。使用聚类从头开始训练模型会导致准确率不佳。\n",
        "- 与前面的层相反，使用更多冗余参数（例如 `tf.keras.layers.Dense` 和 `tf.keras.layers.Conv2D`）来聚类后面的层。\n",
        "- 在微调期间，先冻结前面的层，然后再冻结聚类的层。将冻结层数视为超参数。根据经验，冻结大多数前面的层对于当前的聚类 API 较为理想。\n",
        "- 避免聚类关键层（例如注意力机制）。\n",
        "\n",
        "**更多提示**：`tfmot.clustering.keras.cluster_weights` API 文档提供了有关如何更改每层的聚类配置的详细信息。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqBdl3uJN4r_"
      },
      "outputs": [],
      "source": [
        "# Create a base model\n",
        "base_model = setup_model()\n",
        "base_model.load_weights(pretrained_weights)\n",
        "\n",
        "# Helper function uses `cluster_weights` to make only \n",
        "# the Dense layers train with clustering\n",
        "def apply_clustering_to_dense(layer):\n",
        "  if isinstance(layer, tf.keras.layers.Dense):\n",
        "    return cluster_weights(layer, **clustering_params)\n",
        "  return layer\n",
        "\n",
        "# Use `tf.keras.models.clone_model` to apply `apply_clustering_to_dense` \n",
        "# to the layers of the model.\n",
        "clustered_model = tf.keras.models.clone_model(\n",
        "    base_model,\n",
        "    clone_function=apply_clustering_to_dense,\n",
        ")\n",
        "\n",
        "clustered_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hN0DgpvD5Add"
      },
      "source": [
        "## 为聚类模型设置检查点和进行反序列化"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hfji5KWN6XCF"
      },
      "source": [
        "**您的用例**：仅 HDF5 模型格式需要此代码（HDF5 权重或其他格式不需要）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w7P67mPk6RkQ"
      },
      "outputs": [],
      "source": [
        "# Define the model.\n",
        "base_model = setup_model()\n",
        "base_model.load_weights(pretrained_weights)\n",
        "clustered_model = cluster_weights(base_model, **clustering_params)\n",
        "\n",
        "# Save or checkpoint the model.\n",
        "_, keras_model_file = tempfile.mkstemp('.h5')\n",
        "clustered_model.save(keras_model_file, include_optimizer=True)\n",
        "\n",
        "# `cluster_scope` is needed for deserializing HDF5 models.\n",
        "with tfmot.clustering.keras.cluster_scope():\n",
        "  loaded_model = tf.keras.models.load_model(keras_model_file)\n",
        "\n",
        "loaded_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cUv-scK-N4sN"
      },
      "source": [
        "## 提高聚类模型的准确率"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-fZZopDBN4sO"
      },
      "source": [
        "对于您的特定用例，您可以考虑以下提示：\n",
        "\n",
        "- 形心初始化在最终优化的模型准确率中起到关键作用。通常，线性初始化优于密度和随机初始化，因为它不会丢失较大的权重。但是，对于在具有双峰分布的权重上使用极少簇的情况，已经观察到密度初始化可以提供更出色的准确率。\n",
        "\n",
        "- 微调聚类模型时，将学习率设置为低于训练中使用的学习率。\n",
        "\n",
        "- 有关提高模型准确率的总体思路，请在“定义聚类模型”下查找您的用例对应的提示。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4DXw7YbyN4sP"
      },
      "source": [
        "## 部署"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Y5zLfPzN4sQ"
      },
      "source": [
        "### 导出大小经过压缩的模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wX4OrHD9N4sQ"
      },
      "source": [
        "**常见误区**：`strip_clustering` 和应用标准压缩算法（例如通过 gzip）对于看到聚类压缩的好处必不可少。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZvuiCBsVN4sR"
      },
      "outputs": [],
      "source": [
        "model = setup_model()\n",
        "clustered_model = cluster_weights(model, **clustering_params)\n",
        "\n",
        "clustered_model.compile(\n",
        "    loss=tf.keras.losses.categorical_crossentropy,\n",
        "    optimizer='adam',\n",
        "    metrics=['accuracy']\n",
        ")\n",
        "\n",
        "clustered_model.fit(\n",
        "    x_train,\n",
        "    y_train\n",
        ")\n",
        "\n",
        "final_model = tfmot.clustering.keras.strip_clustering(clustered_model)\n",
        "\n",
        "print(\"final model\")\n",
        "final_model.summary()\n",
        "\n",
        "print(\"\\n\")\n",
        "print(\"Size of gzipped clustered model without stripping: %.2f bytes\" \n",
        "      % (get_gzipped_model_size(clustered_model)))\n",
        "print(\"Size of gzipped clustered model with stripping: %.2f bytes\" \n",
        "      % (get_gzipped_model_size(final_model)))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "clustering_comprehensive_guide.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
